#中文白话文生成
from transformers import BertTokenizer, GPT2LMHeadModel, TextGenerationPipeline

cache_dir = "../../my_model_cache/gpt2-chinese"
# tokenizer = BertTokenizer.from_pretrained("uer/gpt2-chinese-cluecorpussmall", cache_dir=cache_dir)
# model = GPT2LMHeadModel.from_pretrained("uer/gpt2-chinese-cluecorpussmall", cache_dir=cache_dir)

tokenizer = BertTokenizer.from_pretrained(
    cache_dir + r"/models--uer--gpt2-chinese-cluecorpussmall/snapshots/c2c0249d8a2731f269414cc3b22dff021f8e07a3")
model = GPT2LMHeadModel.from_pretrained(
    cache_dir + r"/models--uer--gpt2-chinese-cluecorpussmall/snapshots/c2c0249d8a2731f269414cc3b22dff021f8e07a3")
#device=0 指定当前的推理设备为第一块GPU;如果没有GPU环境，就去掉该参数
text_generator = TextGenerationPipeline(model, tokenizer, device=0)
out = text_generator("在高中时我很喜欢一个女孩，并且和她上了床，十几年都没忘记过，", max_length=100000, do_sample=True)
print(out)
